import numpy as np
from typing_extensions import Iterator
from ...base import DatabaseBuilder


def latlon2pix(lat, lon, radius, depth):
    """
    将经纬度转换为球面坐标，注意单位是弧度
    """
    x = (radius + depth) * np.cos(lat) * np.cos(lon)
    y = (radius + depth) * np.cos(lat) * np.sin(lon)
    z = (radius + depth) * np.sin(lat)
    return x, y, z


class BinadPlainBuilder(DatabaseBuilder):
    """
    基于二元组的不变量库构建基类

    使用普通方式构建的导航数据库，即不作任何删减，只将原始数据按视场大小划分
    划分方式为最高点处反投影对角线，在当前平面上的最大长度，事先计算好。
    """

    def build(self, data: Iterator, max_field=None) -> np.ndarray:
        """
        max_field : 最大视场角范围，单位为m
        """
        max_field = np.inf if max_field is None else max_field
        # 计算真实三维点中心
        x, y, z = latlon2pix(
            data[0] * np.pi / 180, data[1] * np.pi / 180, self.radius, data[3]
        )
        r = data[2]
        # 计算着陆平面
        plane = np.array((x, y, z, np.ones_like(x))).T
        _, _, Vh = np.linalg.svd(plane)
        # 平面方程
        plane = Vh[-1]
        ################
        # 数值不稳定，改为直接计算
        # # 空间中以球面表达的陨石坑矩阵表示, N * 4 * 4
        # Q = np.array(
        #     [
        #         [np.ones_like(x), np.zeros_like(x), np.zeros_like(x), -x],
        #         [np.zeros_like(x), np.ones_like(x), np.zeros_like(x), -y],
        #         [np.zeros_like(x), np.zeros_like(x), np.ones_like(x), -z],
        #         [-x, -y, -z, x * x + y * y + z * z - r * r],
        #     ]
        # ).T
        # # 计算平面上的陨石坑矩阵表示，N * 3 * 3
        # C = cone_plane_cross(Q, plane)
        # 计算平面上的陨石坑的中心
        # center = center_ellipse(C)
        # center = np.array(center).T
        ################
        # 将中心的空间坐标投影为平面坐标：
        # 构建旋转矩阵，使得法向量对齐到z轴
        u = plane[:3] / np.linalg.norm(plane[:3])
        v = np.cross(u, np.array([1, 0, 0]))
        w = np.cross(u, v)
        w = w / np.linalg.norm(w)
        v = np.cross(w, u)
        R = np.column_stack((v, w, u))
        # 旋转中心
        T = -R @ np.array([0, 0, 1])
        # 旋转矩阵
        R = R.T
        # 旋转后的中心
        C = R @ np.array((x, y, z)) + T[:, None]
        center = C[:2].T
        ###############
        C = np.array(
            [
                [np.ones_like(C[0]), np.zeros_like(C[0]), -C[0]],
                [np.zeros_like(C[0]), np.ones_like(C[0]), -C[1]],
                [-C[0], -C[1], C[0] * C[0] + C[1] * C[1] - r * r],
            ]
        ).T
        invariants = []
        index = np.arange(len(center))
        # 根据视场选取中心计算不变量
        for i in range(len(center)):
            # 计算当前中心的视场范围，仅取序号大于当前值的，保证不重复计算
            idx = (np.linalg.norm(center - center[i], axis=1) < max_field) & (index > i)
            # 取出当前视场范围内的椭圆参数
            # 计算不变量
            I1, I2 = self.invariant(C[i, None], C[idx])
            invariants.append((np.ones_like(I1) * i, index[idx], I1, I2))
        invariants = np.concatenate(invariants, axis=1)
        return {
            "catalog": np.block([[x], [y], [z]]).T,
            "temp": np.block([center, r[:, None]]),
            "idx": invariants[:2].astype(np.int32),
            "descriptor": invariants[2:],
            "constants": {"max_field": max_field},
        }

    def invariant(self, Q1, Q2):
        raise NotImplementedError
